# coding: utf8

import six
from itertools import chain

from flask import request, Response
from locust import stats as locust_stats, runners as locust_runners
#from locust import User, task, events
from locust import events
from prometheus_client import Metric, REGISTRY, exposition


# import queue
# import os
# import sys
# import logging
# from locust import HttpUser, task, between, SequentialTaskSet
# from locust_data.maintenanceinfo import except_data, equip_data, handler_data, handler_query, jbcd_data, jbcd_query, \
#     isjbcd_data, zykqr_sept1_data, login_data, tasklist_data, sfsh_sept2_data, czsh_sept3_data, sgqr_sept4_data, \
#     aqjd_sept5_data, gcqr_sept6_data, zpry_sept7_data, tsys_sept8_data, zjdf_sept9_data, getUsers_query
# import json
# from base_request import my_request
# from login_untils import setup_logging
#
# import logging.config
#
# base_path = os.getcwd()
# sys.path.append(base_path)
# base_path_init = os.path.abspath(os.path.join(os.path.dirname('settings.py'), os.path.pardir))
# # 加载初始化日志文件
#
#
# setup_logging(base_path + "/logfile/log.yaml")
# log = logging.getLogger()


# This test_cases adds an external web endpoint to the locust master, and makes it serve as a prometheus exporter.
# Runs it as a normal test_cases, then points prometheus to it.
# locust -f prometheus_exporter.py --master

# Lots of code taken from [mbolek's locust_exporter](https://github.com/mbolek/locust_exporter), thx mbolek!


class LocustCollector(object):
    registry = REGISTRY

    def __init__(self, environment, runner):
        self.environment = environment
        self.runner = runner

    def collect(self):
        # collect metrics only when locust runner is spawning or running.
        runner = self.runner

        if runner and runner.state in (locust_runners.STATE_SPAWNING, locust_runners.STATE_RUNNING):
            stats = []
            for s in chain(locust_stats.sort_stats(runner.stats.entries), [runner.stats.total]):
                stats.append({
                    "method": s.method,
                    "name": s.name,
                    "num_requests": s.num_requests,
                    "num_failures": s.num_failures,
                    "avg_response_time": s.avg_response_time,
                    "min_response_time": s.min_response_time or 0,
                    "max_response_time": s.max_response_time,
                    "current_rps": s.current_rps,
                    "median_response_time": s.median_response_time,
                    "ninetieth_response_time": s.get_response_time_percentile(0.9),
                    # only total stats can use current_response_time, so sad.
                    # "current_response_time_percentile_95": s.get_current_response_time_percentile(0.95),
                    "avg_content_length": s.avg_content_length,
                    "current_fail_per_sec": s.current_fail_per_sec
                })

            # perhaps StatsError.parse_error in e.to_dict only works in python slave, take notices!
            errors = [e.to_dict() for e in six.itervalues(runner.stats.errors)]

            metric = Metric('locust_user_count', 'Swarmed users', 'gauge')
            metric.add_sample('locust_user_count', value=runner.user_count, labels={})
            yield metric

            metric = Metric('locust_errors', 'Locust requests errors', 'gauge')
            for err in errors:
                metric.add_sample('locust_errors', value=err['occurrences'],
                                  labels={'path': err['name'], 'method': err['method'],
                                          'error': err['error']})
            yield metric

            is_distributed = isinstance(runner, locust_runners.MasterRunner)
            if is_distributed:
                metric = Metric('locust_slave_count', 'Locust number of slaves', 'gauge')
                metric.add_sample('locust_slave_count', value=len(runner.clients.values()), labels={})
                yield metric

            metric = Metric('locust_fail_ratio', 'Locust failure ratio', 'gauge')
            metric.add_sample('locust_fail_ratio', value=runner.stats.total.fail_ratio, labels={})
            yield metric

            metric = Metric('locust_state', 'State of the locust swarm', 'gauge')
            metric.add_sample('locust_state', value=1, labels={'state': runner.state})
            yield metric

            stats_metrics = ['avg_content_length', 'avg_response_time', 'current_rps', 'current_fail_per_sec',
                             'max_response_time', 'ninetieth_response_time', 'median_response_time',
                             'min_response_time',
                             'num_failures', 'num_requests']

            for mtr in stats_metrics:
                mtype = 'gauge'
                if mtr in ['num_requests', 'num_failures']:
                    mtype = 'counter'
                metric = Metric('locust_stats_' + mtr, 'Locust stats ' + mtr, mtype)
                for stat in stats:
                    # Aggregated stat's method label is None, so name it as Aggregated
                    # locust has changed name Total to Aggregated since 0.12.1
                    if 'Aggregated' != stat['name']:
                        metric.add_sample('locust_stats_' + mtr, value=stat[mtr],
                                          labels={'path': stat['name'], 'method': stat['method']})
                    else:
                        metric.add_sample('locust_stats_' + mtr, value=stat[mtr],
                                          labels={'path': stat['name'], 'method': 'Aggregated'})
                yield metric


@events.init.add_listener
def locust_init(environment, runner, **kwargs):
    print("locust init event received")
    if environment.web_ui and runner:
        @environment.web_ui.app.route("/export/prometheus")
        def prometheus_exporter():
            registry = REGISTRY
            encoder, content_type = exposition.choose_encoder(request.headers.get('Accept'))
            if 'name[]' in request.args:
                registry = REGISTRY.restricted_registry(request.args.get('name[]'))
            body = encoder(registry)
            return Response(body, content_type=content_type)

        REGISTRY.register(LocustCollector(environment, runner))

#
# class Dummy(User):
#     @task(20)
#     def hello(self):
#         pass
#
# class SpocketTask(SequentialTaskSet):
#     # wait_time = between(1, 5)
#     # 任务初始化
#     def on_start(self):
#         logging.info("任务参数初始化")
#         # 初始化任务参数
#         self.login_data = login_data
#         # 异常上报数据
#         self.except_data = except_data
#         # 查询设备id数据
#         self.equiement_data = equip_data
#         # 隐患指派数据
#         self.hand_data = handler_data
#         # 隐患查询数据
#         self.handl_query = handler_query
#
#         # 编辑检修作业卡数据
#         self.jbcd_data = jbcd_data
#         # 编辑作业卡查询数据
#         self.jbcd_query = jbcd_query
#         # 下发作业卡
#         self.isjbcd_data = isjbcd_data
#         # 查询检修作业人员
#         self.getUsers_query = getUsers_query
#         # 查询任务列表
#         self.tasklist_data = tasklist_data
#         # 作业卡确认
#         self.zykqr_sept1_data = zykqr_sept1_data
#         # 四方确认
#         self.sfsh_sept2_data = sfsh_sept2_data
#         # 厂长审核
#         self.czsh_sept3_data = czsh_sept3_data
#         # 四方施工确认
#         self.sgqr_sept4_data = sgqr_sept4_data
#         # 安全交底
#         self.aqjd_sept5_data = aqjd_sept5_data
#         # 过程确认
#         self.gcqr_sept6_data = gcqr_sept6_data
#         # 摘牌
#         self.zpry_sept7_data = zpry_sept7_data
#         # 调试验收
#         self.tsys_sept8_data = tsys_sept8_data
#         # 总结打分
#         self.zjdf_sept9_data = zjdf_sept9_data
#
#     # 用户登陆
#     @task
#     def login(self):
#         self.login_data['userName'] = self.user.user_userName
#         self.login_data['password'] = "123456"
#         # print(self.login_data)
#         with self.client.request(method='POST', url='/user/login', json=self.login_data, name="登陆",
#                                  catch_response=True) as response:
#             res = json.loads(response.text)
#             if res.get("code") == 0:
#                 response.success()
#                 log.info("接口名称:登陆,请求参数:{}".format(self.login_data))
#                 # print(response.text)
#             else:
#                 response.failure('Failed!')

#
#
#
#
# class SpocketUser(HttpUser):
#     host = 'http://192.168.10.152:10005'  # 设置url根路径
#     # host = 'http://apiygtpm.yong-gang.cn:8260'  # 设置url根路径
#     wait_time = between(1, 5)  # 任务与任务之间的思考时间
#     # weight = 3 多个用户使用该属性进行权重分配
#     # tasks = [SpocketTask,ProcessTask]  # 指定执行任务列表，2个意思有一部分用户走1，一部分走2
#     # weight = 9
#     tasks = [SpocketTask]
#     # tasks = {SpocketTask:1, ProcessTask:3}
#     # user_data = []
#     # with open('locust_data/userdata.csv', 'r') as userfile:
#     #     for line in userfile.readlines():
#     #         user_data.append(line.strip().split(","))
#     # print(user_data)
#     queueData = queue.Queue()  # 队列实例化
#     for count in range(1, 23):  # 循环数据生成
#         data = {
#             "userName": "Test%d" % count,
#             "password": "123456"
#         }
#         queueData.put_nowait(data)  # 循环加入队列《全部》，循环完继续执行 如果用put可能阻塞，在Queue.put()方法加上 block=False 的参数
#
#     # 用户初始化
#     def on_start(self):
#         logging.info("用户开始初始化1")
#         self.login_data = login_data
#         user_info = self.ininUser()
#         # 数据初始化
#         self.user_name = user_info.get("user_name")
#         log.info("用户名称:{}".format(self.user_name))
#         self.user_userName = user_info.get("user_userName")
#         self.header_exception = user_info.get("header_exception")
#
#         self.user_linecode = user_info.get("user_linecode")
#         self.user_lineName = user_info.get("user_lineName")
#         self.user_id = user_info.get("user_id")
#         log.info("用户id:{}".format(self.user_id))
#         self.user_orgcode = user_info.get("user_orgcode")
#         self.operators_list = []
#
#     # 处理用户属性
#     def ininUser(self):
#         try:
#             # userinfo = self.queueData.get()
#             userinfo = self.queueData.get_nowait()
#             self.login_data['userName'] = userinfo["userName"]
#             self.login_data['password'] = userinfo["password"]
#         # print(self.login_data)
#         except self.queue.Empty:
#             # 队列取空后，直接退出
#             print('no data exist')
#             exit(0)
#         else:
#
#             user_message = my_request.run_main("post", "http://192.168.10.152:10005/user/login", self.login_data)
#             try:
#                 user_info = {}
#                 user_info["user_id"] = user_message.get('data').get('id')
#
#                 # 用户的orgcode
#                 user_info["user_orgcode"] = user_message.get('data').get('orgCode')
#
#                 # 用户的name
#                 user_info["user_name"] = user_message.get('data').get('name')
#                 # 用户的code
#                 user_info["user_userName"] = user_message.get('data').get('userName')
#                 # 用户专业code,number
#                 user_info["user_linecode"] = user_message.get('data').get('profession')
#                 # 用户专业name
#                 user_info["user_lineName"] = user_message.get('data').get('professionName')
#                 # print(user_info)
#                 token = user_message.get('data').get('accessToken')
#                 token = 'Bearer ' + token
#                 user_info["header_exception"] = {'Content-Type': 'application/json;charset=UTF-8',
#                                                  'Authorization': token,
#                                                  'OperationUser': user_info.get("user_userName")}
#                 return user_info
#             except Exception as e:
#                 log.info("登陆失败:没有获取到taoken{}".format(e))
#
#
#
#
# if __name__ == "__main__":
#     import os
#
#
#     #启动master命令
#     os.system('locust -f locustfile.py --web-port=8090 --skip-log-setup')
#     #locust --master  -f prometheus_exporter.py --web-port=8090  配置时序数据监听参数
#     #locust -f locustfile.py --master-port=5558 --web-port=8091 --skip-log-setup主机
#     #locust -f locustfile.py  --worker --master-host=192.168.10.158:5558 --skip-log-setup从机